#include <uapi/linux/ptrace.h>
#include <linux/blk-mq.h>
#include <linux/list.h>

// Define a structure for value
struct val_t {
    u64 ts;
    u32 pid;
    u64 data_len;
    char name[TASK_COMM_LEN];
};

// Define arguments for tracepoint
struct tp_args {
    u64 __unused__;
    dev_t dev;
    sector_t sector;
    unsigned int nr_sector;
    unsigned int bytes;
    char rwbs[8];
    char comm[16];
    char cmd[];
};

// Define a hash key structure
struct hash_key {
    dev_t dev;
    u32 rwflag;
    sector_t sector;
};


#ifdef INCLUDE_PATTERN
struct sector_key_t {
    u32 dev_major;
    u32 dev_minor;
};

enum bio_pattern {
    UNKNOWN,
    SEQUENTIAL,
    RANDOM,
};
#endif

// Define a structure for data
struct data_t {
    u32 pid;
    u32 dev;
    u64 rwflag;
    u64 delta;
    u64 qdelta;
    u64 sector;
    u64 len;
#ifdef INCLUDE_PATTERN
    enum bio_pattern pattern;
#endif
    u64 ts;
    char name[TASK_COMM_LEN];

    u64 blk_qts;
    u64 blk_gts;
    u64 blk_its;
    u64 blk_mts;
    u64 blk_dts;
    u64 blk_cts;
    u64 blk_rqts;
    u64 cpu;
    char event;
};

#ifdef INCLUDE_PATTERN
BPF_HASH(last_sectors, struct sector_key_t, u64);
#endif

// Define BPF hash maps for various data associations
BPF_HASH(infobyreq, struct hash_key, struct val_t);
BPF_PERF_OUTPUT(events);

// Define BPF hash maps for block layer timestamps
BPF_HASH(blk_qhash, struct hash_key, struct val_t);
BPF_HASH(blk_ghash, struct hash_key, struct val_t);
BPF_HASH(blk_ihash, struct hash_key, struct val_t);
BPF_HASH(blk_mhash, struct hash_key, struct val_t);
BPF_HASH(blk_rqhash, struct hash_key, struct val_t);
BPF_HASH(blk_dhash, struct hash_key, struct val_t);

// Function to get device number from gendisk structure
static dev_t ddevt(struct gendisk *disk) {
    return (disk->major  << 20) | disk->first_minor;
}

// Function to get read/write flag from command flags
static int get_rwflag(u32 cmd_flags) {
#ifdef REQ_WRITE
    return !!(cmd_flags & REQ_WRITE);
#elif defined(REQ_OP_SHIFT)
    return !!((cmd_flags >> REQ_OP_SHIFT) == REQ_OP_WRITE);
#else
    return !!((cmd_flags & REQ_OP_MASK) == REQ_OP_WRITE);
#endif
}

#define RWBS_LEN	8

// Function to get read/write flag from a string
static int get_rwflag_tp(char *rwbs) {
    for (int i = 0; i < RWBS_LEN; i++) {
        if (rwbs[i] == 'W')
            return 1;
        if (rwbs[i] == '\\0')
            return 0;
    }
    return 0;
}


// trace event Q
int trace_blkq(struct pt_regs *ctx, struct bio *bio)
{
    struct val_t val = {};
    struct data_t data = {};
    if (bpf_get_current_comm(&data.name, sizeof(data.name)) == 0) {
        data.pid = bpf_get_current_pid_tgid() >> 32;
        data.dev = ddevt(bio->bi_disk);
        data.rwflag = get_rwflag(bio->bi_opf);
        data.sector = bio->bi_iter.bi_sector;
        data.len = bio->bi_iter.bi_size;
        data.ts = bpf_ktime_get_ns();
        data.blk_qts = data.ts;
        data.cpu = bpf_get_smp_processor_id();
        data.event = 'Q';

        struct hash_key key = {
            .dev = data.dev,
            .rwflag = data.rwflag,
            .sector = data.sector
        };
        DISK_FILTER
        if(bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
            val.pid = data.pid;
            val.ts = data.ts;
            val.data_len = data.len;
        }
        blk_qhash.update(&key, &val);
        events.perf_submit(ctx, &data, sizeof(data));
    }
    return 0;
}


// trace event G
int trace_blkg(struct pt_regs *ctx, struct request *req) {
    struct hash_key key = {
        .dev = ddevt(req->__RQ_DISK__),
        .rwflag = get_rwflag(req->cmd_flags),
        .sector = req->__sector
    };
    DISK_FILTER
    struct data_t data = {};
     if (bpf_get_current_comm(&data.name, sizeof(data.name)) == 0) {
        data.pid = bpf_get_current_pid_tgid() >> 32;
        data.dev = key.dev;
        data.rwflag = key.rwflag;
        data.sector = key.sector;
        data.len = req->__data_len;
        data.ts = bpf_ktime_get_ns();
        struct val_t *valq;
        valq = blk_qhash.lookup(&key);
        if(valq != 0) {
            data.blk_qts = valq->ts;
        }
        data.blk_gts = data.ts;
        data.cpu = bpf_get_smp_processor_id();
        data.event = 'G';
        struct val_t val = {};
        if(bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
            val.pid = data.pid;
            val.ts = data.ts;
            val.data_len = data.len;
        }
        blk_ghash.update(&key, & val);
        events.perf_submit(ctx, &data, sizeof(data));
    }
    return 0;
}


// trace event I
int trace_blki(struct tp_args *args)
{
    struct hash_key key = {
        .dev = args->dev,
        .rwflag = get_rwflag_tp(args->rwbs),
        .sector = args->sector
    };
    DISK_FILTER
    struct data_t data = {};
    if (bpf_get_current_comm(&data.name, sizeof(data.name)) == 0) {
        data.pid = bpf_get_current_pid_tgid() >> 32;
        data.dev = key.dev;
        data.rwflag = key.rwflag;
        data.sector = key.sector;
        data.len = args->bytes;
        struct val_t *valg;
        valg = blk_ghash.lookup(&key);
        if(valg != 0) {
            data.blk_gts = valg->ts;
        }
        data.ts = bpf_ktime_get_ns();
        data.blk_its = data.ts;
        data.cpu = bpf_get_smp_processor_id();
        data.event = 'I';
        struct val_t val = {};
        if(bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
            val.pid = data.pid;
            val.ts = data.ts;
            val.data_len = data.len;
        }
        blk_ihash.update(&key, &val);
        events.perf_submit((struct pt_regs *)args, &data, sizeof(data));
    }
    return 0;
}

// trace event M
int trace_blkm(struct pt_regs *ctx, struct request_queue *q, struct request *req, struct bio *bio)
{
    struct hash_key key = {
       .dev = ddevt(req->__RQ_DISK__),
       .rwflag = get_rwflag(req->cmd_flags),
       .sector = req->__sector
    };
    DISK_FILTER
    struct data_t data = {};
    if (bpf_get_current_comm(&data.name, sizeof(data.name)) == 0) {
        data.pid = bpf_get_current_pid_tgid() >> 32;
        data.dev = key.dev;
        data.rwflag = key.rwflag;
        data.sector = key.sector;
        data.len = req->__data_len;
        data.ts = bpf_ktime_get_ns();
        struct hash_key bio_key = {
            .dev = ddevt(bio->bi_disk),
            .rwflag = get_rwflag(bio->bi_opf),
            .sector = bio->bi_iter.bi_sector
        };
        struct val_t *valq;
        valq = blk_qhash.lookup(&bio_key);
        if(valq != 0) {
            data.blk_qts = valq->ts;
        }
        data.cpu = bpf_get_smp_processor_id();
        data.event = 'M';
        struct val_t val = {};
        if(bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
            val.pid = data.pid;
            val.ts = data.ts;
            val.data_len = data.len;
        }
        blk_mhash.update(&key, &val);
        events.perf_submit(ctx, &data, sizeof(data));
    }
    return 0;
}


// trace event D
int trace_blkd(struct pt_regs *ctx, struct request *req)
{
    struct hash_key key = {
        .dev = ddevt(req->__RQ_DISK__),
        .rwflag = get_rwflag(req->cmd_flags),
        .sector = req->__sector
    };

    DISK_FILTER
    struct data_t data = {};
    if (bpf_get_current_comm(&data.name, sizeof(data.name)) == 0) {
        data.pid = bpf_get_current_pid_tgid() >> 32;
        data.dev = key.dev;
        data.rwflag = key.rwflag;
        data.sector = key.sector;
        data.len = req->__data_len;
        data.ts = bpf_ktime_get_ns();
        data.blk_dts = data.ts;
        struct val_t *vali;
        vali = blk_ihash.lookup(&key);
        if(vali != 0) {
            data.blk_its = vali->ts;
        }
        struct val_t *valm;
        valm = blk_mhash.lookup(&key);
        if(valm != 0) {
            data.blk_mts = valm->ts;
        }
        data.cpu = bpf_get_smp_processor_id();
        data.event = 'D';
        struct val_t val = {};
        if(bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
            val.pid = data.pid;
            val.ts = data.ts;
            val.data_len = data.len;
        }
        blk_dhash.update(&key, &val);
        events.perf_submit(ctx, &data, sizeof(data));
    }
    return 0;
}


// trace event C
int trace_blkc(struct pt_regs *ctx, struct request *req) {
    struct hash_key key = {
        .dev = ddevt(req->__RQ_DISK__),
        .rwflag = get_rwflag(req->cmd_flags),
        .sector = req->__sector
    };
    DISK_FILTER
    struct data_t data = {};
    if (bpf_get_current_comm(&data.name, sizeof(data.name)) == 0) {
        data.pid = bpf_get_current_pid_tgid() >> 32;
        data.dev = key.dev;
        data.rwflag = key.rwflag;
        data.sector = key.sector;
        data.len = req->__data_len;
        data.ts = bpf_ktime_get_ns();
        struct val_t *valq;
        valq = blk_qhash.lookup(&key);
        if(valq != 0) {
            data.blk_qts = valq->ts;
        }
        struct val_t *vald;
        vald = blk_dhash.lookup(&key);
        if(vald != 0) {
            data.blk_dts = vald->ts;
        }
        data.blk_cts = data.ts;
        data.cpu = bpf_get_smp_processor_id();
        data.event = 'C';
        events.perf_submit(ctx, &data, sizeof(data));
    }
    return 0;
}
